import sys, os, copy, math
from topology_engineer import *
sys.path.append("..")
#import aurora_network
from gurobipy import *

'''		
Design fractional topology based on 
traffic clusters using LP formulation
'''
#combines multiple fractional topology into one
class MultiTrafficTopologyEngineerDifferentRoutingWeightSets(TopologyEngineer):
	def __init__(self, aurora_network, reconfig_length, training_length, all_paths, all_traffic_snapshots, numK, ignore_flows=None):
		num_snapshots = len(all_traffic_snapshots)
		TopologyEngineer.__init__(self, aurora_network, num_snapshots)
		self.flow_pairs_to_ignore = []
		if ignore_flows is not None:
			self.flow_pairs_to_ignore = ignore_flows
		self.toe_class = TopologyEngineeringClass.ROBUST_TOE
		self.training_length = training_length
		self.reconfig_length = reconfig_length
		self.numK = numK
		self.num_snapshots = len(all_traffic_snapshots)
		self.all_traffic_snapshots = all_traffic_snapshots
		self.all_paths = all_paths
		## NOTE: have to rezero the cached logical topologies here because when we initialize the parent class,
		## 		it will assume that the topology is static (see TopologyEngineer class above).
		self.cached_logical_topologies = {}
		return

		# must be called before running other code 
	def cache_logical_topologies(self, round_to_integer=False):
		self.cached_logical_topologies[(0, self.training_length - 1)] = self._generate_uniform_logical_topology()
		## Start computing the logical topologies
		current_time = self.training_length
		print("Caching logical topologies for robust ToE.....")
		while current_time < self.num_snapshots:
			ending_time = min(current_time + self.reconfig_length - 1, self.num_snapshots - 1)
			training_snapshots = self.all_traffic_snapshots[current_time - self.training_length : current_time]
			adj_matrix = self._topology_engineer(training_snapshots, self.all_paths)
			if round_to_integer:
				self.cached_logical_topologies[(current_time, ending_time)] = self.aurora_network.round_fractional_topology_giant_switch(adj_matrix, self.flow_pairs_to_ignore)
			else:
				self.cached_logical_topologies[(current_time, ending_time)] = adj_matrix
			current_time += self.reconfig_length
		print("Caching logical topologies: COMPLETE \n")
		print("first logical topology: \n{}\n".format(self.cached_logical_topologies[(self.training_length, self.training_length + self.reconfig_length - 1)]))
		self.aurora_network.print_stats(self.cached_logical_topologies[(self.training_length, self.training_length + self.reconfig_length - 1)])
		return

	def get_filename_param(self):
		return "robusttoe_r{}t{}k{}".format(self.reconfig_length, self.training_length, self.numK)

	# returns a string description of this class
	def get_string(self):
		return "Robust multi traffic ToE with K clusters : {}".format(self.numK)

	def is_static(self):
		return False

	# given a scale up factor, figure out the routing that minimizes two-hop paths if they exists
	## Todo(jason) : fix this and debug
	def _minimize_multihop(self, traffic_matrix, scaling_factor, all_paths):
		## using QP to reduce multihop reliance
		nblocks = self.aurora_network.get_num_blocks()
		model = Model("minimize multihop")
		model.setParam( 'OutputFlag', False )
		interpod_link_counts = [None] * nblocks
		routing_weights = {}
		link_capacity_constraints = [None] * nblocks
		## setup the link constraints and all optization variables, 
		## which are the link counts and routing variables
		objective_function = QuadExpr()
		for i in range(nblocks):
			interpod_link_counts[i] = [None] * nblocks
			link_capacity_constraints[i] = [None] * nblocks
			for j in range(nblocks):
				if (i != j):
					link_capacity_constraints[i][j] = LinExpr()
					upper_bound = min(self.aurora_network.get_num_links(i), self.aurora_network.get_num_links(j))
					interpod_link_counts[i][j] = model.addVar(lb=0, ub=upper_bound, obj=0, vtype=GRB.CONTINUOUS, name="lc" + str(i) + "_" + str(j))
					weight_index = 0
					for path in all_paths[i][j]:
						var = model.addVar(lb=0, ub=GRB.INFINITY, obj=0., vtype=GRB.CONTINUOUS, name="w_{}".format(path))
						routing_weights[path] = var
						# if multihop path, then add to objective function
						if len(path) > 2: 
							objective_function.add(var * var, mult=1.)

		## add radix degree constraints
		for pod in range(nblocks):
			row_constraint = LinExpr()
			col_constraint = LinExpr()
			nlinks = self.aurora_network.get_num_links(pod)
			for target_pod in range(nblocks):
				if target_pod != pod:
					row_constraint.add(interpod_link_counts[pod][target_pod], mult=1.)
					col_constraint.add(interpod_link_counts[target_pod][pod], mult=1.)
			model.addConstr(lhs=row_constraint, sense=GRB.LESS_EQUAL, rhs=nlinks)
			model.addConstr(lhs=col_constraint, sense=GRB.LESS_EQUAL, rhs=nlinks)

		## add achievability of flows for all i j pairs
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					flow_achievability_constr = LinExpr()
					for path in all_paths[i][j]:
						flow_achievability_constr.add(routing_weights[path], mult=1.)
						path_len = len(path)
						curr_node = path[0]
						for next_node_index in range(1, path_len, 1):
							next_node = path[next_node_index]
							link_capacity_constraints[curr_node][next_node].add(routing_weights[path], mult=1.)
							curr_node = next_node
					model.addConstr(lhs=scaling_factor * traffic_matrix[i][j], sense=GRB.LESS_EQUAL, rhs=flow_achievability_constr)
		
		## add link capacity limit for all paths constraints
		for i in range(nblocks):
			for j in range(nblocks):
				if (i != j):
					capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(i))
					model.addConstr(lhs=link_capacity_constraints[i][j], sense=GRB.LESS_EQUAL, rhs=interpod_link_counts[i][j] * capacity)
		# set up the objective function
		model.setObjective(objective_function, GRB.MINIMIZE)
		# start optimizing
		adj_matrix = np.zeros( (nblocks, nblocks,) )
		try: 
			model.optimize()
			status = model.status
			#print("Status : {}".format(status))
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						adj_matrix[i][j] = interpod_link_counts[i][j].x
			#print adj_matrix
			return adj_matrix
		except GurobiError as e:
			print ("MinimizeMultihop: Error code " + str(e. errno ) + ": " + str(e))
		except AttributeError as e1:
			print ("MinimizeMultihop: Encountered an attribute error : " + str(e1))
	
	def _minimize_mlu_for_all_traffic_matrices(self, traffic_matrices, all_inter_block_paths):
		nblocks = self.aurora_network.get_num_blocks()
		num_tm = len(traffic_matrices)
		model = Model("minimize worst case mlu")
		model.setParam( 'OutputFlag', False )
		## inverse of mlu, i.e. 1 / mlu = beta
		beta = model.addVar(lb=0, ub=GRB.INFINITY, obj=1., vtype=GRB.CONTINUOUS, name="beta")
		# initialize the variables
		# first is fractional topology
		fractional_topology_var = [None] * nblocks
		for i in range(nblocks):
			fractional_topology_var[i] = [None] * nblocks
			for j in range(nblocks):
				if i != j:
					upper_bound = min(self.aurora_network.get_num_links(i), self.aurora_network.get_num_links(j))
					fractional_topology_var[i][j] = model.addVar(lb=0, ub=upper_bound, obj=0, vtype=GRB.CONTINUOUS, name="lc" + str(i) + ":" + str(j))
		# next initialize the m x n x n routing weights vars and link capacity constraints
		# while initilizing the routing weight vars can also enforce the routing weight constraints
		routing_weights_var = [None] * num_tm
		link_capacity_constraints = [None] * num_tm
		for tm, tm_index in zip(traffic_matrices, range(num_tm)):
			routing_weights_var[tm_index] = {}
			link_capacity_constraints[tm_index] = [None] * nblocks
			for i in range(nblocks):
				link_capacity_constraints[tm_index][i] = [None] * nblocks
				for j in range(nblocks):
					if i != j:
						link_capacity_constraints[tm_index][i][j] = LinExpr()
						routing_weights_sum_constraint = LinExpr()
						for path in all_inter_block_paths[i][j]:
							routing_weights_var[tm_index][path] = model.addVar(lb=0., ub=GRB.INFINITY, obj=0, vtype=GRB.CONTINUOUS, name="w{}_{}".format(tm_index, path))
							routing_weights_sum_constraint += routing_weights_var[tm_index][path]
						model.addConstr(lhs=routing_weights_sum_constraint, sense=GRB.EQUAL, rhs=beta)

		
		## Now, begin programming the constraints
		## Radix degree constraints
		for pod in range(nblocks):
			row_constraint = LinExpr()
			col_constraint = LinExpr()
			for target_pod in range(nblocks):
				if target_pod != pod:
					row_constraint.add(fractional_topology_var[pod][target_pod], mult=1.)
					col_constraint.add(fractional_topology_var[target_pod][pod], mult=1.)
			model.addConstr(lhs=row_constraint, sense=GRB.LESS_EQUAL, rhs=float(self.aurora_network.get_num_links(pod)))
			model.addConstr(lhs=col_constraint, sense=GRB.LESS_EQUAL, rhs=float(self.aurora_network.get_num_links(pod)))

		
		## add link capacity limit for all paths constraints
		for tm, tm_index in zip(traffic_matrices, range(num_tm)):
			for i in range(nblocks):
				for j in range(nblocks):
					if (i != j):
						for path in all_inter_block_paths[i][j]:
							src = path[0]
							dst = path[-1]
							curr_node = src
							for next_node in path[1:]:
								link_capacity_constraints[tm_index][curr_node][next_node] += (routing_weights_var[tm_index][path] * tm[src][dst]) 
								curr_node = next_node

		## finally, add in the link utilization constraints
		for tm_index in range(num_tm):
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						capacity = min(self.aurora_network.get_link_capacity(i), self.aurora_network.get_link_capacity(j))
						model.addConstr(lhs=link_capacity_constraints[tm_index][i][j], sense=GRB.LESS_EQUAL, rhs=fractional_topology_var[i][j] * capacity)
		# set up the objective function
		model.setObjective(beta, GRB.MAXIMIZE)
		# start optimizing
		try: 
			model.optimize()
			print("Multi routing set beta : {}".format(beta.x))
			print("Multi routing set ToE has mlu : {}".format(1./beta.x))
			logical_topology_adj_matrix = np.zeros((nblocks, nblocks))
			for i in range(nblocks):
				for j in range(nblocks):
					if i != j:
						logical_topology_adj_matrix[i][j] = fractional_topology_var[i][j].x
			return logical_topology_adj_matrix
		except GurobiError as e:
			print ("MaximizeTrafficScaleUp: Error code " + str(e. errno ) + ": " + str(e))
		except AttributeError :
			print ("MaximizeTrafficScaleUp: Encountered an attribute error ")
		return None


	# conducts topology engineering given the representative traffic matrices
	# aurora network - Aurora Network, contains all the parameters
	# training_traffic_snapshots - all the snapshots used for training the topology
	# number of k - which is the number of traffic clusters we use to solve engineer topology
	def topology_engineer_given_representative_TMs(self, representative_tms, all_inter_block_paths):
		# first get number of blocks and number of links and link capacity
		nblocks = self.aurora_network.get_num_blocks()
		num_representative_tms = len(representative_tms)
		adj_matrix = np.zeros((nblocks, nblocks))
		adj_matrix = self._minimize_mlu_for_all_traffic_matrices(representative_tms, all_inter_block_paths)
		adj_matrix = self.progressive_filling(adj_matrix)
		return adj_matrix, None
		#adj_matrix = self._reinforce_fractional_topology(adj_matrix)
		#return adj_matrix